from ultralytics import YOLO
import cv2
#============ 烟雾明火模型 ============
weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\服务器上的权重\nanshan-fire-v3.pt"
# weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\服务器上的权重\nanshan_fire_20251009.pt"
# weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\服务器上的权重\fire_detect.pt"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火1_1.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火1_2.mp4"
video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火1_3.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火1_4.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火1_5.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火1_6.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火2_1.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火2_2.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火3_1.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\Captures\烟雾明火\烟雾明火3_2.mp4"

# 打开视频文件
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)

# 获取视频的基本信息
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.release()  # 🔥 记得释放

print(f"视频信息: {fps}fps, {width}x{height}, 总帧数:{frame_count}")

model = YOLO(weight_path)
model.predict(
    video_path,
    vid_stride=1,
    imgsz=(height, width),
    classes=[0, 1],           # 🔥 改成 [0, 1] 同时检测烟雾和明火
    save=True,
    show=True,
    show_labels=True,
    show_conf=True,
    conf=0.65,                # 🔥 使用配置文件中的阈值 0.65
    iou=0.3
)

